#ifndef HMS_EIGEN_HELPERS_HPP
#define HMS_EIGEN_HELPERS_HPP

#include "typeAliases.hpp"
#include "typeHelpers.hpp"
#include "unifiedMathFunctions.hpp"

namespace hms
{

/**
 * @brief Use this, when a scalar or a row-wise replicated row vector is required.
 * 
 * @tparam nRows is how many times the row vector is to be replicated.
 * @tparam T 
 * @param arg can be a fundamental type or an Eigen::DenseBase type.
 * @return decltype(auto) 
 */
template<Index nRows, typename T>
HMS_INLINE decltype(auto) repl( const T& arg ){
	using Arg = remove_qualifiers<decltype(arg)>;
	if constexpr ( isDenseType<Arg>() ){
		static_assert( nRows != Eigen::Dynamic );
		static_assert( Arg::RowsAtCompileTime == 1 );
		return arg.template replicate<nRows,1>();
	} else {
		static_assert( std::is_fundamental_v<Arg> );
		return arg;
	}
}

/**
 * @brief Unifies Eigen::select and the conditional operator ?:
 * If ConditionType is an Eigen::DenseBase type, Eigen::select is used,
 * otherwise the conditional operator is used.
 * 
 * @tparam ConditionType 
 * @tparam A 
 * @tparam B 
 * @param cond 
 * @param a 
 * @param b 
 * @return HMS_INLINE 
 */
template<typename ConditionType, typename A, typename B>
HMS_INLINE decltype(auto) select( ConditionType&& cond, A&& a, B&& b ){
	using C = remove_qualifiers<ConditionType>;
	if constexpr ( isDenseType<C>() ){
		return cond.select( std::forward<A>(a), std::forward<B>(b) );
	} else {
		return cond ? std::forward<A>(a) : std::forward<B>(b);
	}
}

namespace internal
{

template<typename A, typename B>
constexpr Index maxRows(){
	if constexpr ( isDenseType<A>() && isDenseType<B>() ){
		return
			A::RowsAtCompileTime > B::RowsAtCompileTime ?
			A::RowsAtCompileTime : B::RowsAtCompileTime;
	} else if constexpr ( isDenseType<A>() ){
		return A::RowsAtCompileTime;
	} else if constexpr ( isDenseType<B>() ){
		return B::RowsAtCompileTime;
	} else {
		return 1;
	}
}

} // namespace internal

template<typename ConditionType, typename A, typename B>
HMS_INLINE decltype(auto) colwiseSelect( ConditionType&& cond, A&& a, B&& b ){
	using C = remove_qualifiers<ConditionType>;
	if constexpr ( isDenseType<C>() ){
		using A_unq = remove_qualifiers<A>;
		using B_unq = remove_qualifiers<B>;
		if constexpr ( isDenseType<A_unq>() )
			assert( cond.cols() == a.cols() );
		if constexpr ( isDenseType<B_unq>() )
			assert( cond.cols() == b.cols() );
		if constexpr ( isDenseType<A_unq>() && isDenseType<B_unq>() ){
			static_assert(
				A_unq::RowsAtCompileTime == B_unq::RowsAtCompileTime
			);
		}
		static constexpr Index replRows { internal::maxRows<A_unq,B_unq>() };
		if constexpr (
			C::RowsAtCompileTime == 1 && (
				isDenseType<A_unq>() ||
				isDenseType<B_unq>()
			) && replRows > 1
		){
			return cond
				.template replicate<replRows, 1>()
				.select(
					std::forward<A>(a), std::forward<B>(b)
				);
		} else {
			return cond.select( std::forward<A>(a), std::forward<B>(b) );
		}
	} else {
		return cond ? std::forward<A>(a) : std::forward<B>(b);
	}
}

} // namespace hms

#endif